from drg_group.chongqing_2024.Base import has_mcc, has_cc, intersect, DF


def M_group(record):
    return record.age >= 60 and len(record.zdList) > 1 and has_mcc(record.zdList[0], record.zdList[1:])


def N_group(record):
    return record.age >= 60 and len(record.zdList) > 1 and (
            has_mcc(record.zdList[0], record.zdList[1:]) or has_cc(record.zdList[0], record.zdList[1:]))


def O_group(record):
    return record.age >= 60 and not (len(record.zdList) > 1 and (
            has_mcc(record.zdList[0], record.zdList[1:]) or has_cc(record.zdList[0], record.zdList[1:])))


def P_group(record):
    return record.age >= 60


def X2_group(record):
    return record.age < 18 and len(record.zdList) > 1 and has_mcc(record.zdList[0], record.zdList[1:])


def X4_group(record):
    return record.age < 18 and len(record.zdList) > 1 and (
            has_mcc(record.zdList[0], record.zdList[1:]) or has_cc(record.zdList[0], record.zdList[1:]))


def X6_group(record):
    return (record.age < 18 and not (len(record.zdList) > 1 and (
            has_mcc(record.zdList[0], record.zdList[1:]) or has_cc(record.zdList[0], record.zdList[1:])))) or (
            record.age < 18 or has_cc(record.zdList[0], record.zdList[1:]))


def X0_group(record):
    return record.age < 18


def X1_group(record):
    return len(record.zdList) > 1 and has_mcc(record.zdList[0], record.zdList[1:])


def X3_group(record):
    return len(record.zdList) > 1 and (
            has_mcc(record.zdList[0], record.zdList[1:]) or has_cc(record.zdList[0], record.zdList[1:]))


def X5_group(record):
    return not (len(record.zdList) > 1 and (
            has_mcc(record.zdList[0], record.zdList[1:]) or has_cc(record.zdList[0], record.zdList[1:]))) or has_cc(
        record.zdList[0], record.zdList[1:])


def X9_group(record):
    return True


def drg_group(adrg, record):
    df = DF
    rows = df[df["ADRG"] == adrg].values.tolist()

    li = []
    for row in rows:
        li.append(row[3][3])

    if 'M' in li:
        if M_group(record):
            return adrg + 'M'
    if 'N' in li:
        if N_group(record):
            return adrg + 'N'
    if 'O' in li:
        if O_group(record):
            return adrg + 'O'
    if 'P' in li:
        if P_group(record):
            return adrg + 'P'
    if '2' in li:
        if X2_group(record):
            return adrg + '2'
    if '4' in li:
        if X4_group(record):
            return adrg + '4'
    if '6' in li:
        if X6_group(record):
            return adrg + '6'
    if '0' in li:
        if X0_group(record):
            return adrg + '0'
    if '1' in li:

        if X1_group(record):
            return adrg + '1'
    if '3' in li:

        if X3_group(record):
            return adrg + '3'
    if '5' in li:

        if X5_group(record):
            return adrg + '5'
    if '9' in li:
        if X9_group(record):
            return adrg + '9'
    return adrg